// FE_solver_2D_Unsteady_NS.cpp : This file contains the 'main' function. Program execution begins and ends there.
//

#include <iostream>
#include "CFEMNonlinear.h"
#include "TimeDiscrete.h"
#define PI 3.14159265
double nu(CPoint p, double t)
{
	return 1.0;
}
double u1_t0(CPoint p)
{
	double x = p.getX();
	double y = p.getY();
	return pow(x, 2)*pow(y, 2) + exp(-y);
}
double u2_t0(CPoint p)
{
	double x = p.getX();
	double y = p.getY();
	return -2.0 / 3.0*x*pow(y, 3) + 2.0 - PI * sin(PI*x);
}
double p_t0(CPoint p)
{
	double x = p.getX();
	double y = p.getY();
	return -(2 - PI * sin(PI*x))*cos(2 * PI*y);
}
double u1_left(CPoint p, double t)
{
	double x = p.getX();
	double y = p.getY();
	return exp(-y)*cos(2 * PI*t);
}
double u1_right(CPoint p, double t)
{
	double x = p.getX();
	double y = p.getY();
	return (pow(y, 2) + exp(-y))*cos(2 * PI*t);
}
double u1_bottom(CPoint p, double t)
{
	double x = p.getX();
	double y = p.getY();
	return (1.0 / 16.0*pow(x, 2) + exp(0.25))*cos(2 * PI*t);
}

double u1_up(CPoint p, double t)
{
	double x = p.getX();
	double y = p.getY();
	return cos(2 * PI*t);
}
double u2_left(CPoint p, double t)
{
	double x = p.getX();
	double y = p.getY();
	return 2 * cos(2 * PI*t);
}
double u2_right(CPoint p, double t)
{
	double x = p.getX();
	double y = p.getY();
	return (-2.0 / 3.0*pow(y, 3) + 2)*cos(2 * PI*t);
}
double u2_bottom(CPoint p, double t)
{
	double x = p.getX();
	double y = p.getY();
	return (1.0 / 96.0*x + 2 - PI * sin(PI*x))*cos(2 * PI*t);
}
double u2_up(CPoint p, double t)
{
	double x = p.getX();
	double y = p.getY();
	return (2 - PI * sin(PI*x))*cos(2 * PI*t);
}

//The function is determined by the analytical solution
//The derivation is completed by the symbol module of python
//see the file :"example1_boundary_condition.py"
double f1(CPoint p, double t)
{
	double x = p.getX();
	double y = p.getY();
	
	return -nu(p,t) * (2 * pow(y,2)* cos(2 * PI*t) + (2 * pow(x,2) + exp(-y))*cos(2 * PI*t)) 
		+ 2 * x*pow(y,2) * (pow(x, 2) * pow(y, 2) + exp(-y))*pow(cos(2 * PI*t),2) 
		+ (2 * pow(x, 2)* y - exp(-y))*(-2 * x*pow(y, 3) / 3.0
		- PI * sin(PI*x) + 2)*pow(cos(2 * PI*t),2) - 2 * PI*(pow(x, 2) * pow(y, 2)
		+ exp(-y))*sin(2 * PI*t) + pow(PI,2) * cos(2 * PI*t)*cos(PI*x)*cos(2 * PI*y);
}

double f2(CPoint p, double t)
{
	double x = p.getX();
	double y = p.getY();	
	return -nu(p, t) * (-4 * x*y*cos(2 * PI*t)
		+ pow(PI, 3) * sin(PI*x)*cos(2 * PI*t))
		- 2 * x*pow(y, 2) * (-2 * x*pow(y, 3) / 3.0
		- PI * sin(PI*x) + 2)*pow(cos(2 * PI*t), 2) 
		+ (-2 * pow(y, 3) / 3
		- pow(PI, 2) * cos(PI*x))*(pow(x, 2) * pow(y, 2)
		+ exp(-y))*pow(cos(2 * PI*t), 2) 
		- 2 * PI*(PI*sin(PI*x) - 2.0)*sin(2 * PI*y)*cos(2 * PI*t)
		- 2 * PI*(-2 * x*pow(y, 3) / 3 - PI * sin(PI*x) + 2)*sin(2 * PI*t);
}
double fun_negative_one(CPoint, double)
{
	return -1.0;
}
double fun_one(CPoint, double)
{
	return 1.0;
}

int main()
{
	int unknown_size = 3;
	CFEMNonlinear fem(unknown_size);
	//read mesh from file
	//the mesh file contains the P,T,Pb,Tb,boundary nodes,boundary edges, and the type of the element
	fem.readMesh("..\\mesh\\cha8_exa1.neu");

	//The initializing process need mesh's information, so the function readMesh should be called first
	fem.Init();

	fem.assignBoundNodesValue(0, 0, Dirichlet, u1_left);
	fem.assignBoundNodesValue(1, 0, Dirichlet, u1_bottom);
	fem.assignBoundNodesValue(2, 0, Dirichlet, u1_right);
	fem.assignBoundNodesValue(3, 0, Dirichlet, u1_up);

	fem.assignBoundNodesValue(0, 1, Dirichlet, u2_left);
	fem.assignBoundNodesValue(1, 1, Dirichlet, u2_bottom);
	fem.assignBoundNodesValue(2, 1, Dirichlet, u2_right);
	fem.assignBoundNodesValue(3, 1, Dirichlet, u2_up);

	CBasisFunction u_basis_trial(TWO_DIM_QUADRATIC);
	CBasisFunction u_basis_test(TWO_DIM_QUADRATIC);

	CBasisFunction p_basis_trial(TWO_DIM_LINEAR);
	CBasisFunction p_basis_test(TWO_DIM_LINEAR);


	pair<int, int> matrix_size_u{ fem.getPbTest().size(),fem.getPbTrial().size() };
	pair<int, int> matrix_size_p{ fem.getPbTest().size(),fem.getP().size() };

	//time discrete
	//the time step should be 8*h^3
	double time_left = 0.0, time_right = 1.0;
	int number_of_time_step = 100;
	double theta = 0.5;
	CTimeDiscrete time(time_left, time_right, number_of_time_step, theta);

	//Maximum number of Newton iteration steps
	int max_iteration = 10;

	//convergence tolerance
	//it means that if the difference between the current iterated step and previous one is less than the tol,
	//then the iteration is convergent even the iteration step doesn't reach the maximum.
	double tol = 1.0e-5;

	//assemble matrix A
	SpMat A1 = fem.assembleMatrix2D(nu, 0, u_basis_trial, 1, 0, fem.getTbTrial(), u_basis_test, 1, 0, fem.getTbTest(), matrix_size_u);
	SpMat A2 = fem.assembleMatrix2D(nu, 0, u_basis_trial, 0, 1, fem.getTbTrial(), u_basis_test, 0, 1, fem.getTbTest(), matrix_size_u);
	SpMat A3 = fem.assembleMatrix2D(nu, 0, u_basis_trial, 1, 0, fem.getTbTrial(), u_basis_test, 0, 1, fem.getTbTest(), matrix_size_u);
	SpMat A5 = fem.assembleMatrix2D(fun_negative_one, 0, p_basis_trial, 0, 0, fem.getT(), u_basis_test, 1, 0, fem.getTbTest(), matrix_size_p);
	SpMat A6 = fem.assembleMatrix2D(fun_negative_one, 0, p_basis_trial, 0, 0, fem.getT(), u_basis_test, 0, 1, fem.getTbTest(), matrix_size_p);
	SpMat zero_1(fem.getP().size(), fem.getP().size());
	zero_1.setZero();
	SpMat A = fem.matrix_vstack(
		fem.matrix_hstack(2 * A1 + A2, A3, A5), 
		fem.matrix_hstack(A3.transpose(), 2 * A2 + A1, A6), 
		fem.matrix_hstack(A5.transpose(), A6.transpose(), zero_1));

	//assemble mass matrix M
	SpMat Me = fem.assembleMatrix2D(fun_one, 0, u_basis_trial, 0, 0, fem.getTbTrial(), u_basis_test, 0, 0, fem.getTbTest(), matrix_size_u);
	SpMat zero_2(matrix_size_p.first, matrix_size_p.second);
	SpMat zero_3(matrix_size_u.first, matrix_size_u.second);
	zero_2.setZero();
	zero_3.setZero();
	SpMat M = fem.matrix_vstack(
		fem.matrix_hstack(Me, zero_3, zero_2), 
		fem.matrix_hstack(zero_3, Me, zero_2), 
		fem.matrix_hstack(zero_2.transpose(), zero_2.transpose(), zero_1));
	//zero vector for the assemble of the right side vector
	//it is not dependent on the time or iteration, so declares in the domain out of the time and iteration loop.
	Rsv zero_vec(fem.getP().size());
	zero_vec.setZero();

	//set the initial value and x0 also represents the solution of previous time step.
	Rsv u1_0 = fem.generate_Initial(u1_t0, fem.getPbTest());
	Rsv u2_0 = fem.generate_Initial(u2_t0, fem.getPbTest());
	Rsv p_0  = fem.generate_Initial(p_t0,  fem.getP());

	//x0 is the solution of current time step.
	Rsv x = fem.vector_stack(u1_0, u2_0, p_0);	
	
	//loop over time step
	for (int m = 0; m < time.getTimeStepNumber(); ++m)
	{
		double dt = time.getTimeStep(m);		
		double tmp1 = time.getTime(m + 1);
		cout << "Time step:\t" << m << ", time:\t" << tmp1 << endl;

		//generate b. b is not dependent on the unknowns and only dependent on the time,
		//so the generation of b is not needed in the iteration loop.		
		Rsv b1 = fem.assembleVector2D(f1, tmp1, u_basis_test, 0, 0, fem.getTbTest());
		Rsv b2 = fem.assembleVector2D(f2, tmp1, u_basis_test, 0, 0, fem.getTbTest());
		
		Rsv b = fem.vector_stack(b1, b2, zero_vec);

		//iterating step
		int l = 1;
		//x_ite_0 is the solution of the previous iteration step.
		//before the iteration ,x_ite_0 is the value of previous time step.
		Rsv x_ite_0 = x;		

		//solution of the current iteration step.
		Rsv x_ite_1;
		while (true)
		{
			cout << "\nIterating step:\t" << l << endl;			
			Rsv u1_ite_0 = x_ite_0.segment(0, matrix_size_u.first);
			Rsv u2_ite_0 = x_ite_0.segment(matrix_size_u.first, matrix_size_u.first);

			//assemble matrix for the nonlinear terms of the NS equations.
			SpMat AN1 = fem.assembleMatrix2DNonlinear(
				u1_ite_0, tmp1, u_basis_trial, 1, 0, fem.getTbTrial(), //coefficient
				u_basis_trial, 0, 0, fem.getTbTrial(), u_basis_test, 0, 0, fem.getTbTest(), matrix_size_u);
			SpMat AN2 = fem.assembleMatrix2DNonlinear(
				u1_ite_0, tmp1, u_basis_trial, 0, 0, fem.getTbTrial(), //coefficient
				u_basis_trial, 1, 0, fem.getTbTrial(), u_basis_test, 0, 0, fem.getTbTest(), matrix_size_u);
			SpMat AN3 = fem.assembleMatrix2DNonlinear(
				u2_ite_0, tmp1, u_basis_trial, 0, 0, fem.getTbTrial(), //coefficient
				u_basis_trial, 0, 1, fem.getTbTrial(), u_basis_test, 0, 0, fem.getTbTest(), matrix_size_u);
			SpMat AN4 = fem.assembleMatrix2DNonlinear(
				u1_ite_0, tmp1, u_basis_trial, 0, 1, fem.getTbTrial(), //coefficient
				u_basis_trial, 0, 0, fem.getTbTrial(), u_basis_test, 0, 0, fem.getTbTest(), matrix_size_u);
			SpMat AN5 = fem.assembleMatrix2DNonlinear(
				u2_ite_0, tmp1, u_basis_trial, 1, 0, fem.getTbTrial(), //coefficient
				u_basis_trial, 0, 0, fem.getTbTrial(), u_basis_test, 0, 0, fem.getTbTest(), matrix_size_u);
			SpMat AN6 = fem.assembleMatrix2DNonlinear(
				u2_ite_0, tmp1, u_basis_trial, 0, 1, fem.getTbTrial(), //coefficient
				u_basis_trial, 0, 0, fem.getTbTrial(), u_basis_test, 0, 0, fem.getTbTest(), matrix_size_u);			
			SpMat AN = fem.matrix_vstack(
				fem.matrix_hstack(AN1 + AN2 + AN3, AN4, zero_2), 
				fem.matrix_hstack(AN5, AN6 + AN2 + AN3, zero_2), 
				fem.matrix_hstack(zero_2.transpose(), zero_2.transpose(), zero_1));

			//assemble vector for the nonlinear terms of NS equations.
			Rsv bN1 = fem.assembleVector2DNonlinear(tmp1,//time
				u1_ite_0, u_basis_trial, 0, 0, fem.getTbTrial(), //coefficient1
				u1_ite_0, u_basis_trial, 1, 0, fem.getTbTrial(), //coefficient2
				u_basis_test, 0, 0, fem.getTbTest());
			Rsv bN2 = fem.assembleVector2DNonlinear(tmp1,//time
				u2_ite_0, u_basis_trial, 0, 0, fem.getTbTrial(),//coefficient1
				u1_ite_0, u_basis_trial, 0, 1, fem.getTbTrial(),//coefficient2
				u_basis_test, 0, 0, fem.getTbTest());
			Rsv bN3 = fem.assembleVector2DNonlinear(tmp1,//time
				u1_ite_0, u_basis_trial, 0, 0, fem.getTbTrial(),//coefficient1
				u2_ite_0, u_basis_trial, 1, 0, fem.getTbTrial(),//coefficient2
				u_basis_test, 0, 0, fem.getTbTest());
			Rsv bN4 = fem.assembleVector2DNonlinear(tmp1,//time
				u2_ite_0, u_basis_trial, 0, 0, fem.getTbTrial(),//coefficient1
				u2_ite_0, u_basis_trial, 0, 1, fem.getTbTrial(),//coefficient2
				u_basis_test, 0, 0, fem.getTbTest());
			Rsv bN = fem.vector_stack(
				bN1 + bN2, 
				bN3 + bN4, 
				zero_vec);

			//the total matrix and vector of the current iteration step
			SpMat Atilde = M / dt + A + AN;
			Rsv btilde = b + M / dt * x + bN;

			//treat boundary conditions
			fem.treatBoundaryConditions(fun_one, tmp1, Atilde, btilde, vector<unsigned int>{fem.getPbTrial().size(), fem.getPbTrial().size(), fem.getP().size()});
			
			double value = -2.0*cos(2 * PI*tmp1);
			fem.imposeFixValue(9 + 2 * matrix_size_u.first, value, Atilde, btilde);
			//fem.PrintMat(Atilde);
			//fem.PrintVec(btilde);
			Atilde.makeCompressed();

			x_ite_1 = fem.solveLinearEqua(DIREC_PardisoLU, Atilde, btilde);	
			
			double residual = fabs(x_ite_1(0) - x_ite_0(0));
			for (size_t i = 0; i < x_ite_1.size(); i++)
			{
				if (residual < fabs(x_ite_1(i) - x_ite_0(i)))
				{
					residual = fabs(x_ite_1(i) - x_ite_0(i));
				}
			}			
			cout << "Max residual of the x=\t" << residual << endl;
			

			if (residual <= tol)
			{
				cout << "\nConvergent, begin next time step" << endl;
				x = x_ite_1;				

				string filename = "u_" + to_string(m+1) + ".vtu";
				//fem.export2VTK(filename, x,fem.getPbTrial(),fem.getTbTrial(),2);
				fem.export2VTKXml(filename, "velocity", tmp1,x, fem.getPbTrial(), fem.getTbTrial(), 2);

				Rsv x_p(fem.getP().size());
				for (int i = 0; i < x_p.rows(); ++i)
				{
					x_p(i) = x(i + 2 * fem.getPbTest().size());
				}
				filename = "p_" + to_string(m+1) + ".vtu";
				fem.export2VTKXml(filename, "pressure", tmp1,x_p, fem.getP(), fem.getT(),1);
				break;//stop the iteration
			}
			else//continue the iteration
			{
				++l;
				if (l == max_iteration)
				{
					cout << "\nStop: Reach the maximum iteration steps and not convergent.Try another initial guess or increase the maximum iteration step" << endl;
					exit(100);
				}
				cout << "\nNot convergent, begin next iteration" << endl;
				x_ite_0 = x_ite_1;//not convergent, assign the value of current iteration step to the previous iteration step and continue iteration.
				
			}

		}
	}

}

